import torch
import torch.nn as nn
import torch.fft

def fftshift3d(img):
    bs, ch, z, h, w = img.shape
    fs11 = img[:, :, :, h//2:, w//2:]
    fs12 = img[:, :, :, h//2:, :w//2]
    fs21 = img[:, :, :, :h//2, w//2:]
    fs22 = img[:, :, :, :h//2, :w//2]
    temp1 = torch.cat([fs11, fs21], axis=3)
    temp2 = torch.cat([fs12, fs22], axis=3)
    output = torch.cat([temp1, temp2], axis=4)
    # output = tf.image.resize_images(output, (size_psc, size_psc), 0)
    return output


def fftshift2d(img, size_psc=128):
    bs, ch, h, w = img.shape
    fs11 = img[:, :, h//2:, w//2:]
    fs12 = img[:, :, h//2:, :w//2]
    fs21 = img[:, :, :h//2, w//2:]
    fs22 = img[:, :, :h//2, :w//2]
    temp1 = torch.cat([fs11, fs21], axis=2)
    temp2 = torch.cat([fs12, fs22], axis=2)
    output = torch.cat([temp1, temp2], axis=3)
    # output = tf.image.resize_images(output, (size_psc, size_psc), 0)
    return output


class RCAB(nn.Module):
    def __init__(self, dim=2, input_channels=64, hidden_channels=[64, 4], output_channels=None):  # size_psc：crop_size input_shape：depth
        super().__init__()
        ConvNd = nn.Conv2d if dim == 2 else nn.Conv3d
        AdaptiveAvgPoolNd = nn.AdaptiveAvgPool2d if dim == 2 else nn.AdaptiveAvgPool3d
        self.fftshiftNd = fftshift2d if dim == 2 else fftshift3d
        self.dim = dim
        output_channels = output_channels or input_channels

        self.conv_gelu1 = nn.Sequential(ConvNd(input_channels, hidden_channels[0], kernel_size=3, stride=1, padding=1),
                                        nn.GELU())
        self.conv_gelu2 = nn.Sequential(ConvNd(hidden_channels[0], hidden_channels[0], kernel_size=3, stride=1, padding=1),
                                        nn.GELU())
        self.conv_gelu_out = nn.Sequential(ConvNd(hidden_channels[0], output_channels, kernel_size=3, stride=1, padding=1),
                                        nn.GELU())

        self.conv_relu1 = nn.Sequential(ConvNd(hidden_channels[0], hidden_channels[0], kernel_size=3, stride=1, padding=1),
                                        nn.ReLU())
        self.avg_pool = AdaptiveAvgPoolNd(1)
        self.conv_relu2 = nn.Sequential(ConvNd(hidden_channels[0], hidden_channels[1], kernel_size=1, stride=1, padding=0),
                                        nn.ReLU())
        self.conv_sigmoid = nn.Sequential(ConvNd(hidden_channels[1], output_channels, kernel_size=1, stride=1, padding=0),
                                          nn.Sigmoid())

    def forward(self, x, gamma=0.8):
        x0 = x
        x = self.conv_gelu1(x)
        x1 = self.conv_gelu_out(x)
        x = self.conv_gelu2(x)
        x = torch.fft.fftn(x, dim=[- i for i in range(1, self.dim + 1)])
        x = torch.pow(torch.abs(x) + 1e-8, gamma)  # abs
        # x = self.fftshiftNd(x)
        x = self.conv_relu1(x)
        x = self.avg_pool(x)
        x = self.conv_relu2(x)
        x = self.conv_sigmoid(x)
        x = x1 * x
        x = x0 + x
        return x


class ResGroup(nn.Module):
    def __init__(self, n_RCAB=4):  # size_psc：crop_size input_shape：depth
        super().__init__()
        RCABs = []
        for _ in range(n_RCAB):
            RCABs.append(RCAB())
        self.RCABs = nn.Sequential(*RCABs)

    def forward(self, x):
        x0 = x
        x = self.RCABs(x)
        x = x0+x
        return x


class DFCAN(nn.Module):
    # size_psc：crop_size input_shape：depth
    def __init__(self, input_shape, scale=2, size_psc=128):
        super().__init__()
        self.input = nn.Sequential(nn.Conv2d(input_shape, 64, kernel_size=3, stride=1, padding=1),
                                   nn.GELU(),)
        n_ResGroup = 4
        ResGroups = []
        for _ in range(n_ResGroup):
            ResGroups.append(ResGroup(n_RCAB=4))
        self.RGs = nn.Sequential(*ResGroups)
        self.conv_gelu = nn.Sequential(nn.Conv2d(64, 64*(scale ** 2), kernel_size=3, stride=1, padding=1),
                                       nn.GELU())
        self.pixel_shuffle = nn.PixelShuffle(scale)
        self.conv_sigmoid = nn.Sequential(
            nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1),)
        # nn.Sigmoid())

    def forward(self, x):
        x = self.input(x)
        x = self.RGs(x)
        x = self.conv_gelu(x)
        x = self.pixel_shuffle(x)  # upsampling
        x = self.conv_sigmoid(x)
        return x


if __name__ == '__main__':
    # #x = Variable(torch.rand(2,1,64,64)).cuda()
    # x = torch.rand(1, 6, 128, 128)

    # #model = UNet().cuda()
    # model = DFCAN(input_shape=x.size()[1])
    # # model.eval()
    # y = model(x)
    # # Output shape: torch.Size([1, 1, 256, 256])
    # print('Output shape:', y.shape)
    # # import hiddenlayer as hl

    x = torch.randn(10, 1, 24, 64, 80)
    model = RCAB(dim=3, input_channels=1)
    y = model(x)
